# -*- coding: utf-8 -*-
# @Time : 2021/02/20
# @Author : ricky
# @File : aescrypt.py
# @Software: vscode
"""
AES加解密
"""
from Crypto.Cipher import AES
import base64

AES_LENGTH = 16


class AesCrypt():
    def __init__(self, key='ngf8t58XzWMYSMLK'):
        self.key = self.add_16(key)
        self.mode = AES.MODE_ECB
        self.aes = AES.new(self.key, self.mode)

    def add_16(self, par):
        if type(par) == str:
            par = par.encode()
        while len(par) % 16 != 0:
            par += b'\x00'
        return par

    def encrypt(self, text):
        """
        加密

        参数：
            text：明文
        """
        text = self.add_16(text)
        self.encrypt_text = self.aes.encrypt(text)
        return str(base64.encodebytes(self.encrypt_text),
                   encoding='utf-8').strip('\0')

    def decrypt(self, text):
        """
        解密

        参数：
            text：密文
        """
        text = base64.decodebytes(text.encode(encoding='utf-8'))
        self.decrypt_text = self.aes.decrypt(text).strip(b"\x00")
        return str(self.decrypt_text, encoding='utf-8')